#!/usr/bin/env python
# -*- coding: utf-8 -*-

import envs
import genpy
import rospy
import actionlib

from arm_object import ArmObject
from robot_object import RobotObject

from trajectory_msgs.msg import JointTrajectoryPoint
from control_msgs.msg import FollowJointTrajectoryAction, FollowJointTrajectoryGoal


class Arm():
    def __init__(self):
        # 机械臂客户端
        self.joint_trajectory_client = actionlib.SimpleActionClient('xarm_controller/follow_joint_trajectory', FollowJointTrajectoryAction)

        # 夹爪客户端
        self.gripper_joint_trajectory_client = actionlib.SimpleActionClient('xarm_gripper_controller/follow_joint_trajectory', FollowJointTrajectoryAction)

        # 机械臂检测物体
        self.arm_object = ArmObject()

        # 小车
        self.robot_object = RobotObject()

        # 执行时间
        self.duration_time = 5

        # 机械臂待机姿态
        self.standby_joint_position = FollowJointTrajectoryGoal()
        self.standby_joint_position_point = JointTrajectoryPoint()
        self.standby_joint_position.trajectory.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        self.standby_joint_position_point.positions = [0.00, -0.67, -1.81, -0.44, 0.00]
        self.standby_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.standby_joint_position.trajectory.points = [self.standby_joint_position_point]

        # 机械臂识别物体姿态
        self.detect_joint_position = FollowJointTrajectoryGoal()
        self.detect_joint_position_point = JointTrajectoryPoint()
        self.detect_joint_position.trajectory.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        self.detect_joint_position_point.positions = [0.00, -0.24, -1.83, -1.56, 0.00]
        self.detect_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.detect_joint_position.trajectory.points = [self.detect_joint_position_point]

        # 机械臂抓取姿态
        self.pick_joint_position = FollowJointTrajectoryGoal()
        self.pick_joint_position_point = JointTrajectoryPoint()
        self.pick_joint_position.trajectory.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        # self.pick_joint_position_point.positions = [0.00, 0.88, -1.83, -0.53, 0.00]  # 不管第一第五关节
        self.pick_joint_position_point.positions = [0.00, 1.35, -0.80, -1.16, 0.00]  # 不管第一第五关节
        self.pick_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.pick_joint_position.trajectory.points = [self.pick_joint_position_point]

        # 机械臂放置姿态
        self.place_joint_position = FollowJointTrajectoryGoal()
        self.place_joint_position_point = JointTrajectoryPoint()
        self.place_joint_position.trajectory.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        self.place_joint_position_point.positions = [0.00, 0.90, -1.50, -0.63, 0.00]
        self.place_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.place_joint_position.trajectory.points = [self.place_joint_position_point]

        # 夹爪待机姿态
        self.gripper_standby_joint_position = FollowJointTrajectoryGoal()
        self.gripper_standby_joint_position_point = JointTrajectoryPoint()
        self.gripper_standby_joint_position.trajectory.joint_names = ['gripper_right_knuckle_joint']
        self.gripper_standby_joint_position_point.positions = [-0.60]
        self.gripper_standby_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.gripper_standby_joint_position.trajectory.points = [self.gripper_standby_joint_position_point]

        # 夹爪抓取姿态
        self.gripper_pick_joint_position = FollowJointTrajectoryGoal()
        self.gripper_pick_joint_position_point = JointTrajectoryPoint()
        self.gripper_pick_joint_position.trajectory.joint_names = ['gripper_right_knuckle_joint']
        self.gripper_pick_joint_position_point.positions = [-0.50]
        self.gripper_pick_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.gripper_pick_joint_position.trajectory.points = [self.gripper_pick_joint_position_point]

        # 夹爪放置姿态
        self.gripper_place_joint_position = FollowJointTrajectoryGoal()
        self.gripper_place_joint_position_point = JointTrajectoryPoint()
        self.gripper_place_joint_position.trajectory.joint_names = ['gripper_right_knuckle_joint']
        self.gripper_place_joint_position_point.positions = [-1.56]
        self.gripper_place_joint_position_point.time_from_start = genpy.Duration(self.duration_time)
        self.gripper_place_joint_position.trajectory.points = [self.gripper_place_joint_position_point]

    def active_callback(self):
        """
        action激活后回调（一次）
        """
        rospy.logdebug('joint trajectory action is actived')

    def done_callback(self, state, res):
        """
        action完成后回调（一次）
        """
        rospy.logdebug('joint trajectory action is finished')
        rospy.logdebug('state: %s', state)
        rospy.logdebug('res: %s', res)

    def feedback_callback(self, feedback):
        """
        action执行过程中反馈的回调（多次）
        """
        rospy.logdebug('joint trajectory action is feedbacking...')
        rospy.logdebug('feedback: %s', feedback)

    def joint_trajectory_send_goal(self, goal):
        """
        action请求（机械臂）
        """
        self.joint_trajectory_client.wait_for_server()

        self.joint_trajectory_client.send_goal(goal, done_cb=self.done_callback, active_cb=self.active_callback, feedback_cb=self.feedback_callback)
        self.joint_trajectory_client.wait_for_result(rospy.Duration.from_sec(300.0))

    def gripper_joint_trajectory_send_goal(self, goal):
        """
        action请求（夹爪）
        """
        self.gripper_joint_trajectory_client.wait_for_server()

        self.gripper_joint_trajectory_client.send_goal(goal, done_cb=self.done_callback, active_cb=self.active_callback, feedback_cb=self.feedback_callback)
        self.gripper_joint_trajectory_client.wait_for_result(rospy.Duration.from_sec(300.0))

    def arm_detect(self):
        """
        机械臂识别垃圾
        """
        rospy.loginfo('[arm] arm detect...')

        # self.gripper_joint_trajectory_send_goal(self.gripper_place_joint_position)
        self.joint_trajectory_send_goal(self.detect_joint_position)

        rospy.loginfo('[arm] arm detect OK')

    def arm_pick(self):
        """
        机械臂抓取垃圾
        """
        rospy.loginfo('[arm] arm pick...')

        self.robot_object.robot_track()
        self.arm_object.arm_swing()
        joint_states = self.arm_object.arm_theta()

        self.gripper_joint_trajectory_send_goal(self.gripper_place_joint_position)

        # 第一关节和第五关节
        self.pick_joint_position_point.positions[0] = joint_states[0]
        self.pick_joint_position_point.positions[4] = joint_states[4]
        self.pick_joint_position.trajectory.points = [self.pick_joint_position_point]
        self.joint_trajectory_send_goal(self.pick_joint_position)

        self.gripper_joint_trajectory_send_goal(self.gripper_pick_joint_position)
        self.joint_trajectory_send_goal(self.standby_joint_position)

        rospy.loginfo('[arm] arm pick OK')

    def arm_place(self):
        """
        机械臂放置垃圾
        """
        rospy.loginfo('[arm] arm place...')

        self.joint_trajectory_send_goal(self.place_joint_position)
        self.gripper_joint_trajectory_send_goal(self.gripper_place_joint_position)
        self.joint_trajectory_send_goal(self.standby_joint_position)
        self.gripper_joint_trajectory_send_goal(self.gripper_standby_joint_position)

        rospy.loginfo('[arm] arm place OK')


if __name__ == '__main__':
    rospy.init_node('npu_robot_arm', anonymous=True, log_level=envs.log_level)
    arm = Arm()
    arm.arm_detect()
    # arm.arm_pick()
    # arm.arm_place()
